package com.sparrow.rpc.serialize.impl;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.sparrow.rpc.serialize.Serializer;
import com.sparrow.rpc.socket.msg.req.RpcLoginRequest;
import com.sparrow.rpc.socket.msg.req.RpcRequest;
import com.sparrow.rpc.socket.msg.resp.RpcResponse;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.util.ArrayList;
import java.util.HashMap;

public class KryoSerializer extends Serializer {

    private Kryo kryo;

    public KryoSerializer() {
        kryo = new Kryo();
        // 不强制注册所有类，避免失败
        kryo.setRegistrationRequired(false);
        register(RpcRequest.class);
        register(RpcResponse.class);
        register(RpcLoginRequest.class);
        register(ArrayList.class);
        register(HashMap.class);
        register(long[].class);
    }

    public void register(Class clazz) {
        kryo.register(clazz);
    }

    @Override
    public <T> byte[] serialize(T obj) {
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
             Output output = new Output(baos)) {
            kryo.writeObject(output, obj);
            output.flush();
            return baos.toByteArray();
        } catch (Exception e) {
            throw new RuntimeException("Serialization error", e);
        }
    }

    @Override
    public <T> T deserialize(byte[] bytes, Class<T> clazz) {
        try (ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes);
             Input input = new Input(byteArrayInputStream)) {
            return kryo.readObject(input, clazz);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}